import os
import sys
import logging
import json
import yaml
from datetime import datetime
from typing import Dict, Any, List, Tuple
import numpy as np
import cv2

# 尝试导入OpenCV的CUDA模块
has_cv2_cuda = False
try:
    cv2.cuda
    has_cv2_cuda = True
except AttributeError:
    pass


def setup_logger(log_file: str = None, log_level: str = 'INFO') -> logging.Logger:
    """
    设置日志记录器
    
    Args:
        log_file: 日志文件路径
        log_level: 日志级别
        
    Returns:
        logging.Logger: 配置好的日志记录器
    """
    # 创建日志记录器
    logger = logging.getLogger('remote_sensing_crop_extraction')
    logger.setLevel(getattr(logging, log_level.upper()))
    
    # 清除已存在的处理器
    if logger.handlers:
        logger.handlers.clear()
    
    # 创建控制台处理器
    console_handler = logging.StreamHandler()
    console_handler.setLevel(getattr(logging, log_level.upper()))
    
    # 创建格式化器
    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    console_handler.setFormatter(formatter)
    
    # 添加处理器到记录器
    logger.addHandler(console_handler)
    
    # 如果指定了日志文件，创建文件处理器
    if log_file:
        # 确保日志目录存在
        log_dir = os.path.dirname(os.path.abspath(log_file))
        os.makedirs(log_dir, exist_ok=True)
        
        file_handler = logging.FileHandler(log_file, encoding='utf-8')
        file_handler.setLevel(getattr(logging, log_level.upper()))
        file_handler.setFormatter(formatter)
        logger.addHandler(file_handler)
    
    return logger


def load_json(file_path: str) -> Dict[str, Any]:
    """
    加载JSON文件
    
    Args:
        file_path: JSON文件路径
        
    Returns:
        Dict[str, Any]: JSON数据
    """
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            return json.load(f)
    except Exception as e:
        print(f"加载JSON文件失败: {e}")
        return {}


def save_json(data: Dict[str, Any], file_path: str) -> bool:
    """
    保存数据到JSON文件
    
    Args:
        data: 要保存的数据
        file_path: JSON文件路径
        
    Returns:
        bool: 保存是否成功
    """
    try:
        # 确保目录存在
        os.makedirs(os.path.dirname(os.path.abspath(file_path)), exist_ok=True)
        
        with open(file_path, 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
        return True
    except Exception as e:
        print(f"保存JSON文件失败: {e}")
        return False


def load_yaml(file_path: str) -> Dict[str, Any]:
    """
    加载YAML文件
    
    Args:
        file_path: YAML文件路径
        
    Returns:
        Dict[str, Any]: YAML数据
    """
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            return yaml.safe_load(f)
    except Exception as e:
        print(f"加载YAML文件失败: {e}")
        return {}


def save_yaml(data: Dict[str, Any], file_path: str) -> bool:
    """
    保存数据到YAML文件
    
    Args:
        data: 要保存的数据
        file_path: YAML文件路径
        
    Returns:
        bool: 保存是否成功
    """
    try:
        # 确保目录存在
        os.makedirs(os.path.dirname(os.path.abspath(file_path)), exist_ok=True)
        
        with open(file_path, 'w', encoding='utf-8') as f:
            yaml.dump(data, f, allow_unicode=True, default_flow_style=False)
        return True
    except Exception as e:
        print(f"保存YAML文件失败: {e}")
        return False


def get_file_list(directory: str, extensions: List[str] = None) -> List[str]:
    """
    获取目录下所有指定扩展名的文件
    
    Args:
        directory: 目录路径
        extensions: 文件扩展名列表，不指定则返回所有文件
        
    Returns:
        List[str]: 文件路径列表
    """
    file_list = []
    
    try:
        for root, _, files in os.walk(directory):
            for file in files:
                if extensions:
                    if any(file.lower().endswith(ext.lower()) for ext in extensions):
                        file_list.append(os.path.join(root, file))
                else:
                    file_list.append(os.path.join(root, file))
    except Exception as e:
        print(f"获取文件列表失败: {e}")
    
    return file_list


def resize_image(image: np.ndarray, max_size: int = 2048, maintain_aspect_ratio: bool = True) -> np.ndarray:
    """
    调整图像大小
    
    Args:
        image: 输入图像
        max_size: 最大尺寸（宽度或高度）
        maintain_aspect_ratio: 是否保持宽高比
        
    Returns:
        np.ndarray: 调整后的图像
    """
    h, w = image.shape[:2]
    
    if maintain_aspect_ratio:
        # 计算缩放比例
        scale = max_size / max(h, w)
        new_h = int(h * scale)
        new_w = int(w * scale)
        
        # 调整大小
        resized_image = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
    else:
        # 直接调整到指定大小
        resized_image = cv2.resize(image, (max_size, max_size), interpolation=cv2.INTER_LINEAR)
    
    return resized_image


def normalize_image(image: np.ndarray) -> np.ndarray:
    """
    归一化图像数据
    
    Args:
        image: 输入图像
        
    Returns:
        np.ndarray: 归一化后的图像
    """
    # 将图像数据转换为浮点型
    normalized_image = image.astype(np.float32)
    
    # 归一化到[0, 1]范围
    min_val = normalized_image.min()
    max_val = normalized_image.max()
    
    if max_val > min_val:
        normalized_image = (normalized_image - min_val) / (max_val - min_val)
    
    return normalized_image


def create_timestamp() -> str:
    """
    创建当前时间的时间戳字符串
    
    Returns:
        str: 时间戳字符串
    """
    return datetime.now().strftime('%Y%m%d_%H%M%S')


def calculate_execution_time(func):
    """
    计算函数执行时间的装饰器
    
    Args:
        func: 要装饰的函数
        
    Returns:
        function: 装饰后的函数
    """
    def wrapper(*args, **kwargs):
        start_time = datetime.now()
        result = func(*args, **kwargs)
        end_time = datetime.now()
        execution_time = (end_time - start_time).total_seconds()
        print(f"函数 {func.__name__} 执行时间: {execution_time:.2f}秒")
        return result
    
    return wrapper


def convert_to_uint8(image: np.ndarray) -> np.ndarray:
    """
    将图像转换为8位无符号整数格式
    
    Args:
        image: 输入图像
        
    Returns:
        np.ndarray: 转换后的图像
    """
    if image.dtype == np.uint8:
        return image
    
    # 归一化到[0, 255]范围
    image_min = image.min()
    image_max = image.max()
    
    if image_max > image_min:
        uint8_image = ((image - image_min) / (image_max - image_min) * 255).astype(np.uint8)
    else:
        uint8_image = np.zeros_like(image, dtype=np.uint8)
    
    return uint8_image


def colorize_mask(mask: np.ndarray, color: Tuple[int, int, int] = (255, 0, 0), alpha: float = 0.5) -> np.ndarray:
    """
    为二值掩码添加颜色和透明度（支持GPU加速）
    
    Args:
        mask: 二值掩码
        color: RGB颜色值
        alpha: 透明度
        
    Returns:
        np.ndarray: 彩色掩码
    """
    # 对于较大的图像使用GPU加速
    if has_cv2_cuda and mask.shape[0] > 100 and mask.shape[1] > 100:
        try:
            # 将掩码转换为float32并上传到GPU
            gpu_mask = cv2.cuda_GpuMat()
            gpu_mask.upload(mask.astype(np.float32))
            
            # 创建彩色掩码（RGBA格式）
            gpu_color_mask = cv2.cuda_GpuMat(mask.shape[0], mask.shape[1], cv2.CV_8UC4)
            
            # 设置颜色通道
            # 创建BGR颜色（注意OpenCV使用BGR格式）
            bgr_color = (color[2], color[1], color[0])  # RGB -> BGR
            
            # 初始化彩色掩码为指定颜色
            gpu_color_mask.create(gpu_color_mask.size(), gpu_color_mask.type(), bgr_color + (0,))
            
            # 计算alpha通道（使用GPU进行乘法运算）
            gpu_alpha = cv2.cuda.multiply(gpu_mask, alpha, dtype=cv2.CV_8U)
            gpu_alpha = cv2.cuda.multiply(gpu_alpha, 255.0, dtype=cv2.CV_8U)
            
            # 将alpha通道复制到彩色掩码的第4通道
            # 由于cv2.cuda没有直接操作多通道的简单方法，我们下载到CPU处理这一步
            color_mask = gpu_color_mask.download()
            color_mask[..., 3] = gpu_alpha.download()
            
            return color_mask
        except Exception as e:
            print(f"GPU着色掩码失败，回退到CPU: {e}")
            # 回退到CPU实现
            pass
    
    # CPU实现（保留原有代码）
    color_mask = np.zeros((mask.shape[0], mask.shape[1], 4), dtype=np.uint8)
    color_mask[..., :3] = color
    color_mask[..., 3] = (mask * 255 * alpha).astype(np.uint8)
    
    return color_mask


def overlay_mask(image: np.ndarray, mask: np.ndarray, color: Tuple[int, int, int] = (255, 0, 0), alpha: float = 0.5) -> np.ndarray:
    """
    将掩码叠加到图像上（支持GPU加速）
    
    Args:
        image: 原始图像
        mask: 二值掩码
        color: RGB颜色值
        alpha: 透明度
        
    Returns:
        np.ndarray: 叠加后的图像
    """
    # 确保图像是3通道RGB格式
    if len(image.shape) == 2:
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
    
    # 对于较大的图像使用GPU加速
    if has_cv2_cuda and image.shape[0] > 100 and image.shape[1] > 100:
        try:
            # 将图像转换为BGR格式（OpenCV默认格式）并上传到GPU
            if image.shape[2] == 3:
                # 如果是RGB格式，转换为BGR
                image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
            else:
                image_bgr = image.copy()
            
            gpu_image = cv2.cuda_GpuMat()
            gpu_image.upload(image_bgr)
            
            # 创建彩色掩码（使用GPU加速的colorize_mask函数）
            color_mask = colorize_mask(mask, color, alpha)
            
            # 将彩色掩码转换为BGR格式并上传到GPU
            color_mask_bgr = cv2.cvtColor(color_mask, cv2.COLOR_RGBA2BGR)
            gpu_color_mask = cv2.cuda_GpuMat()
            gpu_color_mask.upload(color_mask_bgr)
            
            # 使用GPU进行图像叠加
            gpu_overlay = cv2.cuda.addWeighted(gpu_image, 1.0 - alpha, gpu_color_mask, alpha, 0.0)
            
            # 下载结果到CPU
            overlay_bgr = gpu_overlay.download()
            
            # 转换回RGB格式
            if image.shape[2] == 3:
                overlay = cv2.cvtColor(overlay_bgr, cv2.COLOR_BGR2RGB)
            else:
                overlay = overlay_bgr
            
            return overlay
        except Exception as e:
            print(f"GPU叠加掩码失败，回退到CPU: {e}")
            # 回退到CPU实现
            pass
    
    # CPU实现（保留原有代码）
    # 创建彩色掩码
    color_mask = colorize_mask(mask, color, alpha)
    
    # 叠加掩码到图像
    overlay = cv2.addWeighted(image, 1, color_mask[..., :3], alpha, 0)
    
    return overlay